fix(agent): bound code input and tool output byte budgets (#180, #181)#191
Merged
Conversation
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Replace unbounded `wrap_code` with `wrap_code_with_budget` that truncates oversized code to `AgentConfig.max_code_bytes`, using the same UTF-8-safe truncation pattern as `wrap_listing_with_budget`. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
ToolRegistry::execute() now accepts max_output_bytes and truncates its result before returning, so tool output is bounded at the source instead of allocating the full result and truncating after-the-fact. The truncate() helper now reserves marker bytes within the budget so the total output (body + marker) never exceeds the limit. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Each tool now caps its own internal allocations instead of relying solely on the boundary truncate() in execute(): - exec_read_file: uses Read::take() to read only budget+1 bytes from disk, avoiding full-file allocation for large files - grep_recursive: tracks accumulated byte count and stops scanning when byte_budget is reached - list_recursive: same byte accumulator pattern as grep The boundary truncate() in execute() remains as a safety net. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
…ments Bundles results/max_results/total_bytes/byte_budget into a shared struct used by grep_recursive and list_recursive, reducing both from 8/7 args to 4 args. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
read_to_string with Read::take() errors on mid-codepoint truncation. Switch to read_to_end + from_utf8 validation so files with multi-byte characters at the budget boundary are handled gracefully. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
📝 WalkthroughWalkthroughThis PR implements end-to-end byte-budget constraints for code review content. The agent now enforces a configurable ChangesCode and Tool Output Budgeting
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Possibly related issues
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Review rate limit: 9/10 reviews remaining, refill in 6 minutes. Comment |
![coderabbitai[bot]](https://avatars.github.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/agent.rs (1)
218-273:⚠️ Potential issue | 🟠 Major | ⚡ Quick winError outputs bypass remaining-byte cap and accounting
On Line 269 and Line 272, error strings are returned without truncation, and
total_bytes_readis not updated for those branches. Repeated tool/JSON errors can therefore consume prompt space without consuming budget.Proposed fix
- let result = match serde_json::from_str::<serde_json::Value>(&tc.arguments) { + let raw_result = match serde_json::from_str::<serde_json::Value>(&tc.arguments) { Ok(args) => { match tools.execute(&tc.name, &args, remaining) { - Ok(output) => { - // existing truncation/accounting path... - output - } + Ok(output) => output, Err(e) => format!("Error: {}", e), } } Err(e) => format!("Error: malformed arguments: {}", e), }; - Some(result) + let result = truncate(&raw_result, remaining); + if raw_result.len() > remaining { + self.total_bytes_read = config.max_bytes_read; + eprintln!("Agent: byte limit ({}) reached", config.max_bytes_read); + } else { + self.total_bytes_read += result.len(); + } + Some(result)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/agent.rs` around lines 218 - 273, The error branches (Err from tools.execute and malformed args from serde_json::from_str) bypass the remaining-byte cap and don't update self.total_bytes_read; modify the code that produces error messages (the format!("Error: {}", e) and malformed args branch) to run through the same truncation logic used for successful tool output (use remaining, TRUNCATION_MARKER and floor_char_boundary to truncate and decide whether to append the marker), and update self.total_bytes_read exactly as in the truncated_path handling (set to config.max_bytes_read when truncated, otherwise add the final output length) so all output paths respect the byte budget; apply this change around the match arms that produce the result value (the tools.execute Err branch and the serde_json::from_str Err branch).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/tools.rs`:
- Around line 22-25: push currently appends the full entry regardless of
remaining budget, allowing a single large entry to overshoot byte_budget; change
push (and its use of total_bytes/results) to first compute remaining =
byte_budget - *self.total_bytes, return early if remaining <= 0, otherwise
truncate entry to at most remaining - 1 bytes (preserving space for the +1 byte
accounted for), update *self.total_bytes by the actual bytes appended +1, and
push the possibly-truncated string into self.results so traversal stops
respecting the byte budget and avoids spikes from a single large entry.
---
Outside diff comments:
In `@src/agent.rs`:
- Around line 218-273: The error branches (Err from tools.execute and malformed
args from serde_json::from_str) bypass the remaining-byte cap and don't update
self.total_bytes_read; modify the code that produces error messages (the
format!("Error: {}", e) and malformed args branch) to run through the same
truncation logic used for successful tool output (use remaining,
TRUNCATION_MARKER and floor_char_boundary to truncate and decide whether to
append the marker), and update self.total_bytes_read exactly as in the
truncated_path handling (set to config.max_bytes_read when truncated, otherwise
add the final output length) so all output paths respect the byte budget; apply
this change around the match arms that produce the result value (the
tools.execute Err branch and the serde_json::from_str Err branch).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 9c35bccf-de05-4f82-82cb-8b5ccc9d725b
📒 Files selected for processing (2)
src/agent.rssrc/tools.rs
📜 Review details
🧰 Additional context used
📓 Path-based instructions (1)
**/*.rs
📄 CodeRabbit inference engine (CLAUDE.md)
**/*.rs: Do not use emojis in code or output
Support AST-based analysis for Rust files (.rs) including complexity, unsafe code blocks, and unwrap() calls
Use clippy as the linter for Rust files
Include Rust ast-grep rules (5 bundled): block-on-in-async, expect-empty-message, ignored-io-result, silent-error-conversion, string-byte-slice
Files:
src/agent.rssrc/tools.rs
🔇 Additional comments (2)
src/tools.rs (1)
293-304: UTF-8-safe truncation and marker reservation look correctLine 298 and Line 302 protect character boundaries, and marker reservation keeps output within
max.src/agent.rs (1)
172-197:wrap_code_with_budgetcorrectly handles tight budgets and UTF-8 boundariesThe wrapper-overhead guard and boundary-safe truncation are aligned with the listing wrapper behavior.
Comment on lines
+22
to
+25
| fn push(&mut self, entry: String) { | ||
| *self.total_bytes += entry.len() + 1; | ||
| self.results.push(entry); | ||
| } |
There was a problem hiding this comment.
Accumulator can overshoot byte budget before stop condition
push always appends the full entry (Line 23-Line 24), so one large grep/list entry can exceed byte_budget significantly before traversal stops. That weakens the early-bound guarantee and can spike memory on long matches.
Proposed fix
impl AccumulatorState<'_> {
fn push(&mut self, entry: String) {
- *self.total_bytes += entry.len() + 1;
- self.results.push(entry);
+ if self.is_full() {
+ return;
+ }
+ let separator = usize::from(!self.results.is_empty());
+ let remaining = self.byte_budget.saturating_sub(*self.total_bytes);
+ if remaining <= separator {
+ return;
+ }
+
+ let mut entry = entry;
+ let max_entry_bytes = remaining - separator;
+ if entry.len() > max_entry_bytes {
+ let safe_end = entry.floor_char_boundary(max_entry_bytes);
+ entry.truncate(safe_end);
+ }
+
+ *self.total_bytes += entry.len() + separator;
+ self.results.push(entry);
}
}🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/tools.rs` around lines 22 - 25, push currently appends the full entry
regardless of remaining budget, allowing a single large entry to overshoot
byte_budget; change push (and its use of total_bytes/results) to first compute
remaining = byte_budget - *self.total_bytes, return early if remaining <= 0,
otherwise truncate entry to at most remaining - 1 bytes (preserving space for
the +1 byte accounted for), update *self.total_bytes by the actual bytes
appended +1, and push the possibly-truncated string into self.results so
traversal stops respecting the byte budget and avoids spikes from a single large
entry.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
max_code_bytesfield toAgentConfig(100KB default) andwrap_code_with_budget()that truncates with in-prompt marker + stderr warning, mirroring the existingwrap_listing_with_budget()pattern.max_output_bytesparameter toToolRegistry::execute()and pushed the cap into each tool:read_fileusesRead::take()to limit bytes from disk,grep/list_filesuse anAccumulatorStatebyte counter to stop accumulating early.AccumulatorStatestruct to bundle recursive accumulator params, fixing a clippytoo_many_argumentswarning.read_to_stringwithRead::take()errors on mid-codepoint splits; switched toread_to_end+from_utf8validation.Implementation plan
docs/plans/2026-05-02-agent-budget-bounds.md
Test plan
agent_config_has_max_code_bytes_default— new field exists with 100KB defaultwrap_code_with_budget_truncates_oversized_input— budget honored with markerwrap_code_with_budget_passes_small_input_unchanged— no truncation under budgetwrap_code_with_budget_handles_budget_smaller_than_tags— degenerate case returns emptyexecute_read_file_respects_max_output_bytes— tool-level cap enforcedexecute_grep_respects_max_output_bytes— grep respects budgetexecute_with_large_budget_returns_full_output— no regression with large budgetread_file_does_not_allocate_beyond_budget— 1MB file capped at 500 bytesgrep_stops_accumulating_at_byte_budget— early termination on byte limitlist_files_stops_accumulating_at_byte_budget— same for list_filesexecute_tool_call_multi_call_budget_accounting— cross-call budget trackingread_file_handles_utf8_boundary_truncation— multi-byte chars at boundary1758 tests passing, 0 failures. Cargo clippy clean on changed files.
Quorum self-review
4 findings, all pre-existing (complexity in
agent_loop,grep_recursive,list_recursive). No in-branch bugs.🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Bug Fixes